package com.yexin.qingzhou.service;

import com.fasterxml.jackson.databind.JsonNode;
import com.fasterxml.jackson.databind.ObjectMapper;
import com.yexin.qingzhou.entity.ChatMessage;
import org.springframework.beans.factory.annotation.Value;
import org.springframework.http.HttpEntity;
import org.springframework.http.HttpHeaders;
import org.springframework.http.HttpMethod;
import org.springframework.http.MediaType;
import org.springframework.http.ResponseEntity;
import org.springframework.stereotype.Service;
import org.springframework.web.client.RestTemplate;

import java.io.BufferedReader;
import java.io.InputStreamReader;
import java.util.HashMap;
import java.util.List;
import java.util.Map;

@Service
public class DeepSeekService {


    private static final String API_KEY = "sk-c84bffdccc404126b70da5e8f6c3f5ef";
    private static final String BASE_URL = "https://api.deepseek.com/v1/chat/completions";
    
    private final RestTemplate restTemplate;

    public DeepSeekService(RestTemplate restTemplate) {
        this.restTemplate = restTemplate;
    }

    public String chat(List<ChatMessage> messages) {
        HttpHeaders headers = new HttpHeaders();
        headers.setContentType(MediaType.APPLICATION_JSON);
        headers.setBearerAuth(API_KEY);

        Map<String, Object> requestBody = new HashMap<>();
        requestBody.put("model", "deepseek-chat");
        requestBody.put("messages", messages);

        HttpEntity<Map<String, Object>> entity = new HttpEntity<>(requestBody, headers);

        ResponseEntity<Map> response = restTemplate.exchange(
            BASE_URL,
            HttpMethod.POST,
            entity,
            Map.class
        );

        Map responseBody = response.getBody();
        List<Map<String, Object>> choices = (List<Map<String, Object>>) responseBody.get("choices");
        Map<String, Object> message = (Map<String, Object>) choices.get(0).get("message");
        return (String) message.get("content");
    }

    public interface StreamCallback {
        void onMessage(String content);
        void onComplete(String fullResponse);
        void onError(Throwable throwable);
    }

    public void chatStream(List<ChatMessage> messages, StreamCallback callback) {
        HttpHeaders headers = new HttpHeaders();
        headers.setContentType(MediaType.APPLICATION_JSON);
        headers.setBearerAuth(API_KEY);

        Map<String, Object> requestBody = new HashMap<>();
        requestBody.put("model", "deepseek-chat");
        requestBody.put("messages", messages);
        requestBody.put("stream", true);  // 启用流式响应

        HttpEntity<Map<String, Object>> entity = new HttpEntity<>(requestBody, headers);

        StringBuilder fullResponse = new StringBuilder();

        try {
            restTemplate.execute(
                BASE_URL,
                HttpMethod.POST,
                request -> {
                    headers.forEach((key, value) -> request.getHeaders().addAll(key, value));
                    request.getBody().write(new ObjectMapper().writeValueAsBytes(requestBody));
                },
                response -> {
                    try (BufferedReader reader = new BufferedReader(
                            new InputStreamReader(response.getBody()))) {
                        String line;
                        while ((line = reader.readLine()) != null) {
                            if (line.startsWith("data: ")) {
                                String data = line.substring(6);
                                if ("[DONE]".equals(data)) {
                                    callback.onComplete(fullResponse.toString());
                                    return null;
                                }
                                
                                JsonNode node = new ObjectMapper().readTree(data);
                                String content = node.path("choices")
                                    .path(0)
                                    .path("delta")
                                    .path("content")
                                    .asText("");
                                
                                if (!content.isEmpty()) {
                                    fullResponse.append(content);
                                    callback.onMessage(content);
                                }
                            }
                        }
                    }
                    return null;
                }
            );
        } catch (Exception e) {
            callback.onError(e);
        }
    }
} 